import itertools
import random
from utils.prettyrepr import prettify_class

from amaze.tile import Tile
from amaze.terrains import Floor, Wall
from amaze.directions import ORDERED, DX, DY


@prettify_class
class Maze(object):
    __slots__ = ("grid", "width", "height", "time",
                 "all_entities", "active_entities",
                 "__weakref__")

    def __init__(self, width, height, fill=Floor):
        if width % 2 == 0:
            width += 1
        if height % 2 == 0:
            height += 1
        self.grid = [[Tile(self, x, y) for x in xrange(width)] for y in xrange(height)]
        self.width = width
        self.height = height
        self.time = 0.0
        self.all_entities = set()
        self.active_entities = set()
        if fill is not None:
            self.fill(fill)

    def __info__(self):
        return ("w=%d, h=%d, %d entities [%d active]" %
                (self.width, self.height,
                 len(self.all_entities),
                 len(self.active_entities)))

    def __getitem__(self, (x, y)):
        return self.grid[y][x]

    def __setitem__(self, (x, y), terrain):
        self.grid[y][x].terrain = terrain

    @property
    def size(self):
        return (self.width, self.height)

    def neighbor_coords(self, x, y, distance=1):
        for direction in ORDERED:
            nx = x + distance * DX[direction]
            ny = y + distance * DY[direction]
            if 0 <= nx < self.width and 0 <= ny < self.height:
                yield (nx, ny)

    def neighbor_tiles(self, x, y, distance=1):
        for nx, ny in self.neighbor_coords(x, y, distance):
            yield self[nx, ny]

    def neighbor_tile(self, x, y, direction, distance=1):
        nx = x + distance * DX[direction]
        ny = y + distance * DY[direction]
        if 0 <= nx < self.width and 0 <= ny < self.height:
            return self[nx, ny]
        raise IndexError("neighbor unavailable")

    def iter_coords(self):
        return itertools.product(xrange(self.width), xrange(self.height))

    def iter_tiles(self):
        for x, y in self.iter_coords():
            yield self[x, y]

    # --------------------------------------------------------------------------
    # Game management methods
    def init(self):
        """Initialize all active entities in the maze. This should be called exactly once before
        the start of a game."""
        for entity in self.active_entities:
            entity.init()

    def step(self, dt):
        """Call step() and then call post_step() on all active entities in the maze. This allows
        all entities to move first, and only then check for collisions."""
        # Make a list here because during step() new entities may
        # become active or some active entities may become inactive.
        # The newly active entities will only be activated on the next step.
        active_entities = list(self.active_entities)
        for entity in active_entities:
            entity.step(dt)
        self.time += dt
        for entity in active_entities:
            entity.post_step()

    # --------------------------------------------------------------------------
    # Manual generation methods
    def fill(self, terrain=Floor):
        """Fill the maze with the argument terrain."""
        for x, y in self.iter_coords():
            self[x, y] = terrain

    def make_outer_box(self, terrain=Wall):
        """Make a box around the maze using the argument terrain."""
        for x, y in itertools.product((0, self.width-1), xrange(self.height)):
            self[x, y] = terrain
        for x, y in itertools.product(xrange(self.width), (0, self.height-1)):
            self[x, y] = terrain

    # --------------------------------------------------------------------------
    # Automatic generation methods
    def noncell_coords(self):
        for x, y in self.iter_coords():
            if x % 2 == 0 or y % 2 == 0:
                yield (x, y)

    def cell_coords(self):
        return itertools.product(xrange(1, self.width, 2), xrange(1, self.height, 2))

    def generate(self, floor=Floor, wall=Wall, loop_prob=0.01, seed=None):
        if seed is not None:
            random.seed(seed)
        # Create initial setup: the maze is filled with 'floor' terrain and then all non-cell
        # positions are set to 'wall' terrain.
        self.fill(floor)
        for x, y in self.noncell_coords():
            self[x, y] = wall
        # Now we use a randomized depth-first traversal of the cell graph. Cells are the tiles at
        # odd coordinates which initially contain walls between them. When we traverse from one
        # cell to the next, the wall between the cells is cleared. At each step, we randomly choose
        # a non-visited neighbor of our current cell and move into it. If all neighbor cells have
        # been visited, we may create a loop with probability 'loop_prob', and then backtrack.
        unvisited_cells = set(self.cell_coords())
        loops = 0
        while len(unvisited_cells) > 0:
            stack = [unvisited_cells.pop()]
            while len(stack) > 0:
                cx, cy = stack[-1]
                neighbors = list(self.neighbor_coords(cx, cy, 2))
                unvisited_neighbors = filter(unvisited_cells.__contains__, neighbors)
                if len(unvisited_neighbors) > 0:
                    # Choose an unvisited neighbor randomly and move into it.
                    nx, ny = neighbor = random.choice(unvisited_neighbors)
                    unvisited_cells.remove(neighbor)
                    stack.append(neighbor)
                else:
                    # All neighbors have been visited: backtrack and optionally create a loop.
                    stack.pop()
                    if random.random() >= loop_prob:
                        continue
                    nx, ny = random.choice(neighbors)
                    loops += 1
                # Clear the wall between the current cell and the chosen neighbor.
                x = (cx + nx) / 2
                y = (cy + ny) / 2
                self[x, y] = floor
        # Make sure that all the cells are reachable.
        self.check_cell_coverage()
        return loops

    def check_cell_coverage(self):
        """After generate(), make sure that all the cells are reachable from any starting point
        (i.e. the cell graph is not disconnect)."""
        unvisited_cells = set(self.cell_coords())
        if len(unvisited_cells) == 0:
            return
        stack = [unvisited_cells.pop()]
        while len(stack) > 0:
            cx, cy = stack.pop()
            neighbors = []
            for nx, ny in self.neighbor_coords(cx, cy, 2):
                x = (cx + nx) / 2
                y = (cy + ny) / 2
                if self[x, y].terrain is Floor:
                    neighbors.append((nx, ny))
            unvisited_neighbors = filter(unvisited_cells.__contains__, neighbors)
            unvisited_cells.difference_update(unvisited_neighbors)
            stack.extend(unvisited_neighbors)
        if len(unvisited_cells) > 0:
            raise Exception("unreachable cells: %s" % (unvisited_cells,))

    # --------------------------------------------------------------------------
    # Static visualization
    def show(self):
        """Code borrowed from http://en.wikipedia.org/wiki/Maze_generation_algorithm"""
        import matplotlib.pyplot as pyplot

        X = [[bool(tile.terrain is Wall) for tile in row] for row in reversed(self.grid)]
        pyplot.imshow(X, cmap=pyplot.cm.binary, interpolation='nearest')
        pyplot.xticks([])
        pyplot.yticks([])
        pyplot.show()


if __name__ == "__main__":
    m = Maze(50, 25)
    m.generate()
    m.show()
